{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conversational Chess using non-OpenAI clients\n",
    "\n",
    "This notebook provides tips for using non-OpenAI models when using functions/tools.\n",
    "\n",
    "The code is based on [this notebook](../agentchat_nested_chats_chess),\n",
    "which provides a detailed look at nested chats for tool use. Please refer to that\n",
    "notebook for more on nested chats as this will be concentrated on tweaks to\n",
    "improve performance with non-OpenAI models.\n",
    "\n",
    "The notebook represents a chess game between two players with a nested chat to\n",
    "determine the available moves and select a move to make.\n",
    "\n",
    "This game contains a couple of functions/tools that the LLMs must use correctly by the\n",
    "LLMs:\n",
    "- `get_legal_moves` to get a list of current legal moves.\n",
    "- `make_move` to make a move.\n",
    "\n",
    "Two agents will be used to represent the white and black players, each associated with\n",
    "a different LLM cloud provider and model:\n",
    "- Anthropic's Sonnet 3.5 will be Player_White\n",
    "- Mistral's Mixtral 8x7B (using Together.AI) will be Player_Black\n",
    "\n",
    "As this involves function calling, we use larger, more capable, models from these providers.\n",
    "\n",
    "The nested chat will be supported be a board proxy agent who is set up to execute\n",
    "the tools and manage the game.\n",
    "\n",
    "Tips to improve performance with these non-OpenAI models will be noted throughout **in bold**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "First, you need to install the `ag2` and `chess` packages to use AG2. We'll include Anthropic and Together.AI libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -qqq ag2[anthropic,together] chess"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up LLMs\n",
    "\n",
    "We'll use the Anthropic (`api_type` is `anthropic`) and Together.AI (`api_type` is `together`) client classes, with their respective models, which both support function calling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Annotated\n",
    "\n",
    "import chess\n",
    "import chess.svg\n",
    "from IPython.display import display\n",
    "\n",
    "from autogen import ConversableAgent, register_function\n",
    "\n",
    "# Let's set our two player configs, specifying clients and models\n",
    "\n",
    "# Anthropic's Sonnet for player white\n",
    "player_white_config_list = [\n",
    "    {\n",
    "        \"api_type\": \"anthropic\",\n",
    "        \"model\": \"claude-3-5-sonnet-20240620\",\n",
    "        \"api_key\": os.getenv(\"ANTHROPIC_API_KEY\"),\n",
    "        \"cache_seed\": None,\n",
    "    },\n",
    "]\n",
    "\n",
    "# Mistral's Mixtral 8x7B for player black (through Together.AI)\n",
    "player_black_config_list = [\n",
    "    {\n",
    "        \"api_type\": \"together\",\n",
    "        \"model\": \"mistralai/Mixtral-8x7B-Instruct-v0.1\",\n",
    "        \"api_key\": os.environ.get(\"TOGETHER_API_KEY\"),\n",
    "        \"cache_seed\": None,\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll setup game variables and the two functions for getting the available moves and then making a move."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the board.\n",
    "board = chess.Board()\n",
    "\n",
    "# Keep track of whether a move has been made.\n",
    "made_move = False\n",
    "\n",
    "\n",
    "def get_legal_moves() -> Annotated[\n",
    "    str,\n",
    "    \"Call this tool to list of all legal chess moves on the board, output is a list in UCI format, e.g. e2e4,e7e5,e7e8q.\",\n",
    "]:\n",
    "    return \"Possible moves are: \" + \",\".join([str(move) for move in board.legal_moves])\n",
    "\n",
    "\n",
    "def make_move(\n",
    "    move: Annotated[\n",
    "        str,\n",
    "        \"Call this tool to make a move after you have the list of legal moves and want to make a move. Takes UCI format, e.g. e2e4 or e7e5 or e7e8q.\",\n",
    "    ],\n",
    ") -> Annotated[str, \"Result of the move.\"]:\n",
    "    move = chess.Move.from_uci(move)\n",
    "    board.push_uci(str(move))\n",
    "    global made_move\n",
    "    made_move = True\n",
    "    # Display the board.\n",
    "    display(\n",
    "        chess.svg.board(board, arrows=[(move.from_square, move.to_square)], fill={move.from_square: \"gray\"}, size=200)\n",
    "    )\n",
    "    # Get the piece name.\n",
    "    piece = board.piece_at(move.to_square)\n",
    "    piece_symbol = piece.unicode_symbol()\n",
    "    piece_name = (\n",
    "        chess.piece_name(piece.piece_type).capitalize()\n",
    "        if piece_symbol.isupper()\n",
    "        else chess.piece_name(piece.piece_type)\n",
    "    )\n",
    "    return f\"Moved {piece_name} ({piece_symbol}) from {chess.SQUARE_NAMES[move.from_square]} to {chess.SQUARE_NAMES[move.to_square]}.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating agents\n",
    "\n",
    "Our main player agents are created next, with a few tweaks to help our models play:\n",
    "\n",
    "- Explicitly **telling agents their names** (as the name field isn't sent to the LLM).\n",
    "- Providing simple instructions on the **order of functions** (not all models will need it).\n",
    "- Asking the LLM to **include their name in the response** so the message content will include their names, helping the LLM understand who has made which moves.\n",
    "- Ensure **no spaces are in the agent names** so that their name is distinguishable in the conversation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "player_white = ConversableAgent(\n",
    "    name=\"Player_White\",\n",
    "    system_message=\"You are a chess player and you play as white, your name is 'Player_White'. \"\n",
    "    \"First call the function get_legal_moves() to get list of legal moves. \"\n",
    "    \"Then call the function make_move(move) to make a move. \"\n",
    "    \"Then tell Player_Black you have made your move and it is their turn. \"\n",
    "    \"Make sure you tell Player_Black you are Player_White.\",\n",
    "    llm_config={\"config_list\": player_white_config_list, \"cache_seed\": None},\n",
    ")\n",
    "\n",
    "player_black = ConversableAgent(\n",
    "    name=\"Player_Black\",\n",
    "    system_message=\"You are a chess player and you play as black, your name is 'Player_Black'. \"\n",
    "    \"First call the function get_legal_moves() to get list of legal moves. \"\n",
    "    \"Then call the function make_move(move) to make a move. \"\n",
    "    \"Then tell Player_White you have made your move and it is their turn. \"\n",
    "    \"Make sure you tell Player_White you are Player_Black.\",\n",
    "    llm_config={\"config_list\": player_black_config_list, \"cache_seed\": None},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create a proxy agent that will be used to move the pieces on the board."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if the player has made a move, and reset the flag if move is made.\n",
    "def check_made_move(msg):\n",
    "    global made_move\n",
    "    if made_move:\n",
    "        made_move = False\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "board_proxy = ConversableAgent(\n",
    "    name=\"Board_Proxy\",\n",
    "    llm_config=False,\n",
    "    # The board proxy will only terminate the conversation if the player has made a move.\n",
    "    is_termination_msg=check_made_move,\n",
    "    # The auto reply message is set to keep the player agent retrying until a move is made.\n",
    "    default_auto_reply=\"Please make a move.\",\n",
    "    human_input_mode=\"NEVER\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our functions are then assigned to the agents so they can be passed to the LLM to choose from.\n",
    "\n",
    "We have tweaked the descriptions to provide **more guidance on when** to use it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_function(\n",
    "    make_move,\n",
    "    caller=player_white,\n",
    "    executor=board_proxy,\n",
    "    name=\"make_move\",\n",
    "    description=\"Call this tool to make a move after you have the list of legal moves.\",\n",
    ")\n",
    "\n",
    "register_function(\n",
    "    get_legal_moves,\n",
    "    caller=player_white,\n",
    "    executor=board_proxy,\n",
    "    name=\"get_legal_moves\",\n",
    "    description=\"Call this to get a legal moves before making a move.\",\n",
    ")\n",
    "\n",
    "register_function(\n",
    "    make_move,\n",
    "    caller=player_black,\n",
    "    executor=board_proxy,\n",
    "    name=\"make_move\",\n",
    "    description=\"Call this tool to make a move after you have the list of legal moves.\",\n",
    ")\n",
    "\n",
    "register_function(\n",
    "    get_legal_moves,\n",
    "    caller=player_black,\n",
    "    executor=board_proxy,\n",
    "    name=\"get_legal_moves\",\n",
    "    description=\"Call this to get a legal moves before making a move.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Almost there, we now create nested chats between players and the board proxy agent to work out the available moves and make the move."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "player_white.register_nested_chats(\n",
    "    trigger=player_black,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            # The initial message is the one received by the player agent from\n",
    "            # the other player agent.\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_white,\n",
    "            # The final message is sent to the player agent.\n",
    "            \"summary_method\": \"last_msg\",\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "\n",
    "player_black.register_nested_chats(\n",
    "    trigger=player_white,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            # The initial message is the one received by the player agent from\n",
    "            # the other player agent.\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_black,\n",
    "            # The final message is sent to the player agent.\n",
    "            \"summary_method\": \"last_msg\",\n",
    "        }\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Playing the game\n",
    "\n",
    "Now the game can begin!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clear the board.\n",
    "board = chess.Board()\n",
    "\n",
    "chat_result = player_black.initiate_chat(\n",
    "    player_white,\n",
    "    message=\"Let's play chess! Your move.\",\n",
    "    max_turns=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this stage, it's hard to tell who's going to win, but they're playing well and using the functions correctly."
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "LLM-backed agents playing chess with each other using nested chats.",
   "tags": [
    "nested chat",
    "tool/function",
    "orchestration"
   ]
  },
  "kernelspec": {
   "display_name": "autogen",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
